import pika
from retry import retry
from pyrabbit.api import Client
from feapder.utils.log import log
from crawler.utils.Singleton import Singleton

from crawler.connection.setting import RABBIT_MQ_HOST, RABBIT_MQ_PORT, RABBIT_MQ_ACCOUNT, RABBIT_MQ_PASSWORD, \
    RABBIT_MQ_MONITOR_PORT

connection_params = {
    'host': RABBIT_MQ_HOST,
    'port': RABBIT_MQ_PORT,
    'credentials': pika.PlainCredentials(RABBIT_MQ_ACCOUNT, RABBIT_MQ_PASSWORD),
    'heartbeat': 600,
    'socket_timeout': 60,
    'blocked_connection_timeout': 300,
    'connection_attempts': 10000,
    'retry_delay': 10
}

connection = None


class RabbitMQConnection:

    # 获取连接
    def _get_connection(self):
        return pika.BlockingConnection(pika.ConnectionParameters(**connection_params))

    # 获取通道，当connection为None时，重新获取
    def _get_channel(self):
        global connection
        try:
            if not connection:
                connection = self._get_connection()
            channel = connection.channel()
        except:
            connection = self._get_connection()
            channel = connection.channel()
        return channel

    # 获取通道
    def get_channel(self):
        global connection
        try:
            channel = connection.channel()
            # 解决channel无法回收的问题
            if channel.channel_number > 100:
                connection.close()
                connection = None
        finally:
            return self._get_channel()

    # 创建队列
    @classmethod
    def create_queue(cls, channel, queue_name):
        return channel.queue_declare(queue=queue_name, durable=True, arguments={'x-max-priority': 10})

    # 获取任务信息
    def get_task(self, queue_name, channel=None):
        if not channel:
            channel = self.get_channel()
        method_frame, header_frame, origin_data = channel.basic_get(queue=queue_name, auto_ack=True)

        if not origin_data:
            # 无任务返回空
            return
        return origin_data, method_frame.message_count

    # 数据插入
    def publish_insert_task(self, task, routing_key, channel=None):
        if not channel:
            channel = self.get_channel()
        channel.basic_publish(
            exchange='',
            routing_key=routing_key,
            body=task,
        )
        log.info(task)
        return channel


@Singleton
class MonitorRabbit:
    def __init__(self, queue_name):
        self.queue_name = queue_name

    @retry(delay=3)
    def get_queue_count(self):
        cl = Client(f'{RABBIT_MQ_HOST}:{RABBIT_MQ_MONITOR_PORT}', RABBIT_MQ_ACCOUNT, RABBIT_MQ_PASSWORD)

        res = cl.get_queue('/', self.queue_name)
        return res['messages']

#
# mr = MonitorRabbit(queue_name='TestCreateTask')
# print(mr.get_queue_count())
